#!/usr/bin/env python3
"""
数据源管理器
支持多个数据源的切换和管理
"""

import abc
from typing import Dict, List, Any, Optional, Tuple
import pandas as pd
from datetime import datetime, timedelta
from enum import Enum

from .utils.logger import get_logger
from .data_sources import BaseDataSource, DataSourceStatus, TushareSource, EastMoneySource, SinaSource, EFinanceSource

logger = get_logger("DataSourceManager")


class DataSourceType(Enum):
    """数据源类型"""
    BAOSTOCK = "baostock"
    YFINANCE = "yfinance"
    AKSHARE = "akshare"
    TUSHARE = "tushare"
    EASTMONEY = "eastmoney"
    SINA = "sina"
    EFINANCE = "efinance"


# 基础类已从 data_sources.base 导入


class BaoStockSource(BaseDataSource):
    """BaoStock数据源"""

    def __init__(self, config: Dict[str, Any] = None):
        super().__init__("baostock", config)
        self.fetcher = None
        self._initialize_fetcher()

    def _initialize_fetcher(self):
        """初始化数据获取器"""
        try:
            from .ashare_fetcher import AShareFetcher
            from .ashare_list import AShareListFetcher
            self.fetcher = AShareFetcher()
            self.list_fetcher = AShareListFetcher()
            self.update_status(DataSourceStatus.AVAILABLE)
            logger.info(f"BaoStock数据源初始化成功")
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            logger.error(f"BaoStock数据源初始化失败: {e}")

    def check_availability(self) -> bool:
        """检查数据源是否可用"""
        try:
            if not self.fetcher:
                self._initialize_fetcher()

            # 简单检查，不进行网络请求
            available = self.fetcher is not None
            self.update_status(
                DataSourceStatus.AVAILABLE if available else DataSourceStatus.ERROR,
                None if available else "BaoStock未可用"
            )
            return available
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            return False

    def get_stock_list(self) -> pd.DataFrame:
        """获取股票列表"""
        try:
            df = self.list_fetcher.get_all_stocks()
            # 过滤退市股票
            df = df[~df['name'].str.contains('退', na=False)]
            # 过滤可转债，只保留股票
            if 'type' in df.columns:
                df = df[df['type'] == '1']
            return df
        except Exception as e:
            logger.error(f"BaoStock获取股票列表失败: {e}")
            return pd.DataFrame()

    def fetch_history_data(self, symbol: str, interval: str,
                          start_date: str, end_date: str) -> pd.DataFrame:
        """获取历史数据"""
        try:
            return self.fetcher.fetch_history_data(symbol, interval, start_date, end_date)
        except Exception as e:
            logger.error(f"BaoStock获取历史数据失败 {symbol}: {e}")
            return pd.DataFrame()

    def get_stock_info(self, symbol: str) -> Dict[str, Any]:
        """获取股票信息"""
        try:
            df = self.list_fetcher.get_all_stocks()
            stock_info = df[df['symbol'] == symbol]
            if not stock_info.empty:
                return stock_info.iloc[0].to_dict()
            return {}
        except Exception as e:
            logger.error(f"BaoStock获取股票信息失败 {symbol}: {e}")
            return {}


class YahooFinanceSource(BaseDataSource):
    """Yahoo Finance数据源"""

    def __init__(self, config: Dict[str, Any] = None):
        super().__init__("yfinance", config)
        self.fetcher = None
        self._initialize_fetcher()

    def _initialize_fetcher(self):
        """初始化数据获取器"""
        try:
            import yfinance as yf
            self.fetcher = yf
            self.update_status(DataSourceStatus.AVAILABLE)
            logger.info(f"Yahoo Finance数据源初始化成功")
        except ImportError:
            self.update_status(DataSourceStatus.ERROR, "yfinance未安装")
            logger.error("yfinance未安装，请运行: pip install yfinance")
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            logger.error(f"Yahoo Finance数据源初始化失败: {e}")

    def check_availability(self) -> bool:
        """检查数据源是否可用"""
        try:
            if not self.fetcher:
                self.update_status(DataSourceStatus.ERROR, "yfinance未初始化")
                return False

            # 简单检查，不进行网络请求
            available = self.fetcher is not None
            self.update_status(
                DataSourceStatus.AVAILABLE if available else DataSourceStatus.ERROR,
                None if available else "Yahoo Finance未可用"
            )
            return available
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            return False

    def get_stock_list(self) -> pd.DataFrame:
        """获取股票列表（Yahoo Finance不提供完整列表）"""
        # Yahoo Finance不提供完整的A股列表，返回空DataFrame
        logger.warning("Yahoo Finance不提供完整的A股列表")
        return pd.DataFrame()

    def fetch_history_data(self, symbol: str, interval: str,
                          start_date: str, end_date: str) -> pd.DataFrame:
        """获取历史数据"""
        try:
            # 转换股票代码格式
            if symbol.startswith('000') or symbol.startswith('001') or symbol.startswith('002'):
                yf_symbol = f"{symbol}.SZ"
            elif symbol.startswith('300'):
                yf_symbol = f"{symbol}.SZ"
            elif symbol.startswith('6'):
                yf_symbol = f"{symbol}.SS"
            else:
                yf_symbol = symbol

            ticker = self.fetcher.Ticker(yf_symbol)

            # 转换间隔格式
            interval_map = {
                '1d': '1d',
                '1h': '1h',
                '30m': '30m',
                '15m': '15m',
                '5m': '5m'
            }
            yf_interval = interval_map.get(interval, '1d')

            data = ticker.history(start=start_date, end=end_date, interval=yf_interval)

            if not data.empty:
                # 转换为标准格式
                data = data.reset_index()
                data = data.rename(columns={
                    'Date': 'Date',
                    'Open': 'Open',
                    'High': 'High',
                    'Low': 'Low',
                    'Close': 'Close',
                    'Volume': 'Volume'
                })
                # 保留需要的列
                data = data[['Date', 'Open', 'High', 'Low', 'Close', 'Volume']]

            return data
        except Exception as e:
            logger.error(f"Yahoo Finance获取历史数据失败 {symbol}: {e}")
            return pd.DataFrame()

    def get_stock_info(self, symbol: str) -> Dict[str, Any]:
        """获取股票信息"""
        try:
            # Yahoo Finance获取详细信息有限
            return {
                'symbol': symbol,
                'source': 'yfinance',
                'available': True
            }
        except Exception as e:
            logger.error(f"Yahoo Finance获取股票信息失败 {symbol}: {e}")
            return {}


class AKShareSource(BaseDataSource):
    """AKShare数据源"""

    def __init__(self, config: Dict[str, Any] = None):
        super().__init__("akshare", config)
        self.fetcher = None
        self._initialize_fetcher()

    def _initialize_fetcher(self):
        """初始化数据获取器"""
        try:
            import akshare as ak
            self.fetcher = ak
            self.update_status(DataSourceStatus.AVAILABLE)
            logger.info(f"AKShare数据源初始化成功")
        except ImportError:
            self.update_status(DataSourceStatus.ERROR, "akshare未安装")
            logger.error("akshare未安装，请运行: pip install akshare")
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            logger.error(f"AKShare数据源初始化失败: {e}")

    def check_availability(self) -> bool:
        """检查数据源是否可用"""
        try:
            if not self.fetcher:
                self.update_status(DataSourceStatus.ERROR, "akshare未初始化")
                return False

            # 简单检查，不进行网络请求
            available = self.fetcher is not None
            self.update_status(
                DataSourceStatus.AVAILABLE if available else DataSourceStatus.ERROR,
                None if available else "AKShare未可用"
            )
            return available
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            return False

    def get_stock_list(self) -> pd.DataFrame:
        """获取股票列表"""
        try:
            df = self.fetcher.stock_zh_a_spot()
            # 过滤退市股票
            df = df[~df['名称'].str.contains('退', na=False)]
            # 重命名列以匹配标准格式
            df = df.rename(columns={
                '代码': 'symbol',
                '名称': 'name',
                '市场类型': 'market_type'
            })
            # 添加其他必要字段
            df['status'] = '0'
            df['type'] = '1'
            return df
        except Exception as e:
            logger.error(f"AKShare获取股票列表失败: {e}")
            return pd.DataFrame()

    def fetch_history_data(self, symbol: str, interval: str,
                          start_date: str, end_date: str) -> pd.DataFrame:
        """获取历史数据"""
        try:
            # 转换间隔格式
            interval_map = {
                '1d': 'daily',
                '1h': '60',
                '30m': '30',
                '15m': '15',
                '5m': '5'
            }
            ak_interval = interval_map.get(interval, 'daily')

            if ak_interval == 'daily':
                data = self.fetcher.stock_zh_a_hist(symbol=symbol, start_date=start_date, end_date=end_date)
            else:
                # 分钟级数据
                data = self.fetcher.stock_zh_a_hist_min_em(symbol=symbol, start_date=start_date, end_date=end_date,
                                                       period=ak_interval, adjust='')

            if not data.empty:
                # 转换为标准格式
                data = data.rename(columns={
                    '日期': 'Date',
                    '开盘': 'Open',
                    '最高': 'High',
                    '最低': 'Low',
                    '收盘': 'Close',
                    '成交量': 'Volume'
                })
                # 保留需要的列
                data = data[['Date', 'Open', 'High', 'Low', 'Close', 'Volume']]

            return data
        except Exception as e:
            logger.error(f"AKShare获取历史数据失败 {symbol}: {e}")
            return pd.DataFrame()

    def get_stock_info(self, symbol: str) -> Dict[str, Any]:
        """获取股票信息"""
        try:
            stock_list = self.fetcher.stock_zh_a_spot()
            stock_info = stock_list[stock_list['代码'] == symbol]
            if not stock_info.empty:
                return stock_info.iloc[0].to_dict()
            return {}
        except Exception as e:
            logger.error(f"AKShare获取股票信息失败 {symbol}: {e}")
            return {}


class DataSourceManager:
    """数据源管理器"""

    def __init__(self, auto_check: bool = False):
        self.sources: Dict[str, BaseDataSource] = {}
        self.active_source: Optional[str] = None
        self._initialize_sources(auto_check)

    def _initialize_sources(self, auto_check: bool = False):
        """初始化所有数据源"""
        sources = [
            # 重新启用数据源，但不进行自动网络检查
            (DataSourceType.BAOSTOCK.value, BaoStockSource()),
            (DataSourceType.YFINANCE.value, YahooFinanceSource()),
            (DataSourceType.AKSHARE.value, AKShareSource()),
            (DataSourceType.TUSHARE.value, TushareSource()),
            (DataSourceType.EASTMONEY.value, EastMoneySource()),
            (DataSourceType.SINA.value, SinaSource()),
            (DataSourceType.EFINANCE.value, EFinanceSource()),
        ]

        for name, source in sources:
            self.sources[name] = source
            if auto_check:
                # 检查数据源可用性
                source.check_availability()
                logger.info(f"数据源 {name}: {source.status.value}")

        if auto_check:
            # 设置默认数据源 - 优先选择可用的数据源
            self._set_best_available_source()

    def check_all_sources(self):
        """手动检查所有数据源的可用性"""
        for name, source in self.sources.items():
            source.check_availability()
            logger.info(f"数据源 {name}: {source.status.value}")

        # 设置默认数据源
        self._set_best_available_source()

    def set_active_source(self, source_name: str) -> bool:
        """设置活跃数据源"""
        if source_name not in self.sources:
            logger.error(f"未知的数据源: {source_name}")
            return False

        source = self.sources[source_name]
        if source.status == DataSourceStatus.AVAILABLE:
            self.active_source = source_name
            logger.info(f"已切换到数据源: {source_name}")
            return True
        else:
            logger.warning(f"数据源 {source_name} 不可用: {source.status.value}")
            return False

    def get_active_source(self) -> Optional[BaseDataSource]:
        """获取当前活跃数据源"""
        if self.active_source:
            return self.sources[self.active_source]
        return None

    def get_available_sources(self) -> List[Dict[str, Any]]:
        """获取所有可用数据源的状态"""
        return [
            {
                'name': name,
                'status': source.status.value,
                'is_active': name == self.active_source,
                'last_check': source.last_check_time,
                'error_message': source.last_error
            }
            for name, source in self.sources.items()
        ]

    def get_combined_stock_list(self) -> pd.DataFrame:
        """获取合并的股票列表"""
        all_stocks = []

        for name, source in self.sources.items():
            if source.status == DataSourceStatus.AVAILABLE:
                try:
                    stocks = source.get_stock_list()
                    if not stocks.empty:
                        stocks['data_source'] = name
                        all_stocks.append(stocks)
                except Exception as e:
                    logger.error(f"获取数据源 {name} 的股票列表失败: {e}")

        if all_stocks:
            # 合并所有股票列表，去重
            combined_df = pd.concat(all_stocks, ignore_index=True)
            # 按symbol去重，优先保留第一个数据源的数据
            combined_df = combined_df.drop_duplicates(subset=['symbol'], keep='first')
            return combined_df

        return pd.DataFrame()

    def fetch_data_with_fallback(self, symbol: str, interval: str,
                                 start_date: str, end_date: str) -> Tuple[pd.DataFrame, str]:
        """使用回退机制获取数据"""
        # 先尝试活跃数据源
        if self.active_source:
            source = self.sources[self.active_source]
            try:
                data = source.fetch_history_data(symbol, interval, start_date, end_date)
                if not data.empty:
                    return data, self.active_source
            except Exception as e:
                logger.error(f"活跃数据源 {self.active_source} 获取数据失败: {e}")

        # 尝试其他可用数据源
        for name, source in self.sources.items():
            if name != self.active_source and source.status == DataSourceStatus.AVAILABLE:
                try:
                    data = source.fetch_history_data(symbol, interval, start_date, end_date)
                    if not data.empty:
                        logger.info(f"使用备用数据源 {name} 获取数据")
                        return data, name
                except Exception as e:
                    logger.error(f"数据源 {name} 获取数据失败: {e}")

        return pd.DataFrame(), ""

    def _set_best_available_source(self):
        """设置最佳可用数据源"""
        # 数据源优先级（根据稳定性和数据质量排序）
        priority_order = [
            DataSourceType.BAOSTOCK.value,
            DataSourceType.EFINANCE.value,
            DataSourceType.TUSHARE.value,
            DataSourceType.EASTMONEY.value,
            DataSourceType.AKSHARE.value,
            DataSourceType.SINA.value,
            DataSourceType.YFINANCE.value,
        ]

        for source_name in priority_order:
            if source_name in self.sources:
                source = self.sources[source_name]
                if source.status == DataSourceStatus.AVAILABLE:
                    self.active_source = source_name
                    logger.info(f"已设置最佳数据源: {source_name}")
                    return

        # 如果没有可用数据源，设置第一个
        if self.sources:
            first_source = list(self.sources.keys())[0]
            self.active_source = first_source
            logger.warning(f"没有可用数据源，使用默认: {first_source}")

    def get_data_source_quality_score(self, source_name: str) -> float:
        """获取数据源质量评分"""
        if source_name not in self.sources:
            return 0.0

        source = self.sources[source_name]

        # 基础分数
        base_score = 0.0
        if source.status == DataSourceStatus.AVAILABLE:
            base_score = 100.0
        elif source.status == DataSourceStatus.LIMITED:
            base_score = 60.0
        elif source.status == DataSourceStatus.UNAVAILABLE:
            base_score = 20.0
        else:  # ERROR
            base_score = 0.0

        # 根据数据源特性调整分数
        quality_adjustments = {
            'baostock': 0.95,      # 专业A股数据，质量高但速度一般
            'efinance': 0.90,      # 新兴库，功能全面
            'tushare': 0.85,       # 老牌数据源，但免费版限制多
            'eastmoney': 0.80,     # 东财数据，质量不错
            'akshare': 0.75,       # 数据全面但稳定性一般
            'sina': 0.70,          # 新浪数据，实时性好但历史数据有限
            'yfinance': 0.60,      # 主要针对美股，A股支持有限
        }

        adjustment = quality_adjustments.get(source_name, 0.5)
        return base_score * adjustment

    def get_best_sources_for_task(self, task_type: str = 'history') -> list:
        """根据任务类型获取最佳数据源列表"""
        available_sources = [
            name for name, source in self.sources.items()
            if source.status == DataSourceStatus.AVAILABLE
        ]

        # 根据任务类型排序
        if task_type == 'history':
            # 历史数据获取优先级
            priority = ['baostock', 'efinance', 'tushare', 'eastmoney', 'akshare']
        elif task_type == 'realtime':
            # 实时数据获取优先级
            priority = ['sina', 'eastmoney', 'efinance', 'akshare', 'baostock']
        elif task_type == 'stock_list':
            # 股票列表获取优先级
            priority = ['baostock', 'eastmoney', 'efinance', 'akshare', 'tushare']
        else:
            priority = ['baostock', 'efinance', 'tushare', 'eastmoney', 'akshare', 'sina']

        # 按优先级排序可用数据源
        sorted_sources = []
        for source_name in priority:
            if source_name in available_sources:
                sorted_sources.append(source_name)

        # 添加其他可用数据源
        for source_name in available_sources:
            if source_name not in sorted_sources:
                sorted_sources.append(source_name)

        return sorted_sources